from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import backtrader as bt
import pandas as pd
import numpy as np
import os
import time
import datetime
import mplfinance as mpf

# ==============================================================================
#  1. 自定义佣金方案
# ==============================================================================
class StampDutyCommissionScheme(bt.CommInfoBase):
    params = (('stamp_duty', 0.0005), ('commission', 0.00025), ('min_commission', 5))
    def _getcommission(self, size, price, pseudoexec):
        comm = abs(size) * price * self.p.commission
        final_comm = max(comm, self.p.min_commission)
        if size < 0: final_comm += abs(size) * price * self.p.stamp_duty
        return final_comm

# ==============================================================================
#  2. 自定义仓位管理器
# ==============================================================================
class RiskParitySizer(bt.Sizer):
    params = (('risk_per_trade_percent', 0.015), ('atr_stop_multiplier', 2.5))
    def _getsizing(self, comminfo, cash, data, isbuy):
        if not isbuy: return self.broker.getposition(data).size
        stock_name = data._name
        atr_indicator = self.strategy.atr.get(stock_name)
        if atr_indicator is None: return 0
        atr = atr_indicator[0]
        if atr <= 0: return 0
        price = data.close[0]; risk_per_share = atr * self.p.atr_stop_multiplier
        if risk_per_share == 0: return 0
        portfolio_value = self.broker.getvalue()
        position_risk_value = portfolio_value * self.p.risk_per_trade_percent
        size = int(position_risk_value / risk_per_share)
        size = (size // 100) * 100
        if size * price > cash: return 0
        return size

# ==============================================================================
#  3. 策略逻辑
# ==============================================================================
class SystematicRiskParityStrategyBT(bt.Strategy):
    params = (('sma_long_window', 150), ('rsi_window', 14), ('rsi_pullback_level', 50),
              ('atr_window', 20), ('trailing_stop_multiplier', 3.0), ('trailing_stop_activation_multiplier', 1.0),
              ('max_holding_days', 30))

    def __init__(self):
        self.sma = {d._name: bt.indicators.SimpleMovingAverage(d.close, period=self.p.sma_long_window) for d in self.datas}
        self.rsi = {d._name: bt.indicators.RSI(d.close, period=self.p.rsi_window) for d in self.datas}
        self.atr = {d._name: bt.indicators.ATR(d, period=self.p.atr_window) for d in self.datas}
        self.pos_data = {}
        self.transaction_log = []

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]: return
        if order.status in [order.Completed]:
            stock_name = order.data._name
            self.transaction_log.append({'date': bt.num2date(order.executed.dt), 'stock': stock_name, 'type': 'BUY' if order.isbuy() else 'SELL', 'price': order.executed.price})
            if order.isbuy():
                self.pos_data[stock_name] = {
                    'entry_price': order.executed.price, 'entry_date': bt.num2date(order.executed.dt).date(), 
                    'entry_atr': self.atr[stock_name][0],
                    'initial_stop_loss': order.executed.price - self.sizer.p.atr_stop_multiplier * self.atr[stock_name][0],
                    'trailing_stop_loss': order.executed.price - self.sizer.p.atr_stop_multiplier * self.atr[stock_name][0]}
                self.log(f'BUY EXECUTED, {stock_name}, Price: {order.executed.price:.2f}, Size: {order.executed.size}')
            elif order.issell(): self.log(f'SELL EXECUTED, {stock_name}, Price: {order.executed.price:.2f}, Size: {order.executed.size}')
        elif order.status in [order.Canceled, order.Margin, order.Rejected]: self.log(f'Order Canceled/Margin/Rejected for {order.data._name}')

    def notify_trade(self, trade):
        if trade.isclosed:
            stock_name = trade.data._name
            if stock_name in self.pos_data: del self.pos_data[stock_name]
            self.log(f'TRADE PROFIT on {stock_name}, NET: {trade.pnlcomm:.2f}')

    def next(self):
        active_positions = len(self.pos_data)
        for d in self.datas:
            stock_name, pos = d._name, self.getposition(d)
            if pos.size > 0:
                pos_info = self.pos_data.get(stock_name)
                if not pos_info: continue
                if d.close[0] < self.sma[stock_name][0]: self.close(data=d); self.log(f'CLOSE: Trend Reversal for {stock_name}'); continue
                if d.close[0] > pos_info['entry_price'] + self.p.trailing_stop_activation_multiplier * pos_info['entry_atr']:
                    new_trailing_stop = d.high[0] - self.p.trailing_stop_multiplier * self.atr[stock_name][0]
                    pos_info['trailing_stop_loss'] = max(pos_info['trailing_stop_loss'], new_trailing_stop)
                stop_price = max(pos_info['initial_stop_loss'], pos_info['trailing_stop_loss'])
                if d.close[0] <= stop_price:
                    reason = "Trailing" if pos_info['trailing_stop_loss'] > pos_info['initial_stop_loss'] else "Initial"
                    self.close(data=d); self.log(f'CLOSE: {reason} Stop for {stock_name}'); continue
                days_held = (bt.num2date(d.datetime[0]).date() - pos_info['entry_date']).days
                if days_held > self.p.max_holding_days and d.close[0] <= pos_info['entry_price']:
                    self.close(data=d); self.log(f'CLOSE: Time Stop for {stock_name}'); continue
            else:
                if active_positions >= 5: continue
                if d.close[0] > self.sma[stock_name][0] and self.rsi[stock_name][0] < self.p.rsi_pullback_level:
                    self.buy(data=d); self.log(f'BUY CREATE, {stock_name}, Price: {d.close[0]:.2f}')

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        print(f'{dt.isoformat()} - {txt}')

# ==============================================================================
#  4. 自定义数据加载器 & 报告生成
# ==============================================================================
class GenericCSV_X(bt.feeds.GenericCSVData):
    params = (('dtformat', ('%Y-%m-%d')), ('datetime', 0), ('open', 1), ('high', 2), ('low', 3), ('close', 4), ('volume', 5), ('openinterest', -1))

def generate_backtrader_report(cerebro, results, output_path):
    print("\n--- Generating Backtest Report ---"); os.makedirs(output_path, exist_ok=True)
    report_path = os.path.join(output_path, 'report.txt'); strategy = results[0]; analyzers = strategy.analyzers
    trade_analysis = analyzers.tradeanalyzer.get_analysis()
    total_trades = trade_analysis.get('total', {}).get('total', 0); won_trades = trade_analysis.get('won', {}).get('total', 0); lost_trades = trade_analysis.get('lost', {}).get('total', 0)
    win_rate = (won_trades / total_trades) * 100 if total_trades > 0 else 0
    total_win_pnl = trade_analysis.get('won', {}).get('pnl', {}).get('total', 0); total_lost_pnl = trade_analysis.get('lost', {}).get('pnl', {}).get('total', 0)
    profit_factor = abs(total_win_pnl / total_lost_pnl) if total_lost_pnl != 0 else float('inf')
    avg_win = total_win_pnl / won_trades if won_trades > 0 else 0; avg_loss = total_lost_pnl / lost_trades if lost_trades > 0 else 0
    pnl_per_trade = (total_win_pnl + total_lost_pnl) / total_trades if total_trades > 0 else 0
    with open(report_path, 'w', encoding='utf-8') as f:
        f.write("="*57 + "\n" + "         Backtrader Backtest Analysis Report\n" + "="*57 + "\n\n")
        f.write(f"Run ID: {os.path.basename(output_path)}\nReport Generated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
        f.write("--- Strategy Parameters ---\n"); [f.write(f"{key:<25}: {value}\n") for key, value in strategy.params._getitems()]; f.write("\n")
        initial_value, final_value = cerebro.broker.startingcash, cerebro.broker.getvalue()
        f.write("--- Overall Performance ---\n"); f.write(f"{'Initial Value':<25}: {initial_value:,.2f}\n"); f.write(f"{'Final Value':<25}: {final_value:,.2f}\n"); f.write(f"{'Total Return':<25}: {(final_value / initial_value) - 1:.2%}\n\n")
        sharpe = analyzers.sharperatio.get_analysis().get('sharperatio', 0.0); drawdown = analyzers.drawdown.get_analysis().get('max', {}).get('drawdown', 0.0)
        f.write("--- Risk & Return Metrics ---\n"); f.write(f"{'Sharpe Ratio (Annualized)':<25}: {sharpe:.3f}\n"); f.write(f"{'Max Drawdown':<25}: {drawdown:.2%}\n\n")
        f.write("--- Trade Analysis Details ---\n"); f.write(f"{'Total Trades':<25}: {total_trades}\n"); f.write(f"{'Winning Trades':<25}: {won_trades}\n"); f.write(f"{'Losing Trades':<25}: {lost_trades}\n")
        f.write(f"{'Win Rate':<25}: {win_rate:.2f} %\n"); f.write(f"{'Profit Factor':<25}: {profit_factor:.2f}\n"); f.write(f"{'Avg. PnL per Trade':<25}: {pnl_per_trade:,.2f}\n")
        f.write(f"{'Avg. Win':<25}: {avg_win:,.2f}\n"); f.write(f"{'Avg. Loss':<25}: {avg_loss:,.2f}\n")
        total_comm = sum(t.commission for t in trade_analysis.get('total', {}).get('trades', [])); f.write(f"{'Total Commission & Fees':<25}: {total_comm:,.2f}\n\n")
    print(f"[Analysis] Detailed report saved to: {report_path}")
    
    charts_path = os.path.join(output_path, 'trade_charts'); os.makedirs(charts_path, exist_ok=True)
    print(f"[Analysis] Generating candlestick dashboards in: '{charts_path}'")
    
    trans_df = pd.DataFrame(strategy.transaction_log)
    if not trans_df.empty: trans_df['date'] = pd.to_datetime(trans_df['date'])

    for i, d in enumerate(cerebro.datas):
        stock_name = d._name
        df = pd.DataFrame({'open': d.open.get(size=len(d)), 'high': d.high.get(size=len(d)), 'low': d.low.get(size=len(d)), 'close': d.close.get(size=len(d)), 'volume': d.volume.get(size=len(d))}, index=[bt.num2date(dt) for dt in d.datetime.get(size=len(d))])
        if df.empty: continue
        
        p = strategy.params
        sma_series = pd.Series(strategy.sma[stock_name].lines.sma.get(size=len(d)), index=df.index)
        rsi_series = pd.Series(strategy.rsi[stock_name].lines.rsi.get(size=len(d)), index=df.index)
        atr_series = pd.Series(strategy.atr[stock_name].lines.atr.get(size=len(d)), index=df.index)

        add_plots = [mpf.make_addplot(sma_series, panel=0, color='blue', width=0.7), mpf.make_addplot(rsi_series, panel=2, color='purple', ylabel='RSI'), mpf.make_addplot(atr_series, panel=3, color='orange', ylabel='ATR')]
        
        plot_title = f'Analysis Dashboard for {stock_name}'
        if not trans_df.empty and stock_name in trans_df['stock'].values:
            plot_title = f'Trades on {stock_name}'
            stock_trans = trans_df[trans_df['stock'] == stock_name]; buy_signals = stock_trans[stock_trans['type'] == 'BUY']; sell_signals = stock_trans[stock_trans['type'] == 'SELL']
            
            # ######################################################
            # ##                核心修改在这里                     ##
            # ######################################################
            if not buy_signals.empty:
                buy_signal_series = pd.Series(buy_signals['price'].values * 0.98, index=buy_signals['date'])
                buy_markers = buy_signal_series.reindex(df.index)
                add_plots.append(mpf.make_addplot(buy_markers, type='scatter', panel=0, marker='^', color='lime', markersize=200))
            if not sell_signals.empty:
                sell_signal_series = pd.Series(sell_signals['price'].values * 1.02, index=sell_signals['date'])
                sell_markers = sell_signal_series.reindex(df.index)
                add_plots.append(mpf.make_addplot(sell_markers, type='scatter', panel=0, marker='v', color='red', markersize=200))
            # ######################################################
        
        chart_file = os.path.join(charts_path, f"dashboard_{stock_name}.png")
        mpf.plot(df, type='candle', style='charles', title=plot_title, ylabel='Price', volume=True, addplot=add_plots,
                 panel_ratios=(8, 1, 2, 2), hlines=dict(hlines=[p.rsi_pullback_level, 70], colors=['g', 'r'], linestyle='--'),
                 savefig=dict(fname=chart_file, dpi=300), figsize=(48, 12), warn_too_much_data=len(df) + 1)
        print(f"  - Dashboard for {stock_name} saved.")

# ==============================================================================
#  5. 主执行逻辑 (Main)
# ==============================================================================
if __name__ == '__main__':
    cerebro = bt.Cerebro()
    data_path = './data'
    universe = [
        "000651.SZ",
        "600000.SH",
        "600036.SH",
        "600519.SH",
    ]
    print(f"Loading {len(universe)} data feeds...");
    for stock in universe:
        dataname = os.path.join(data_path, f'{stock}.csv')
        if os.path.exists(dataname):
            feed = GenericCSV_X(dataname=dataname, fromdate=datetime.datetime(2020, 1, 1), todate=datetime.datetime(2024, 12, 31))
            cerebro.adddata(feed, name=stock)

    cerebro.addstrategy(SystematicRiskParityStrategyBT)
    cerebro.addsizer(RiskParitySizer)
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharperatio')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
    cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='tradeanalyzer')
    
    cerebro.broker.setcash(1000000.0)
    comm_scheme = StampDutyCommissionScheme()
    cerebro.broker.addcommissioninfo(comm_scheme)

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    results = cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    run_id = f"run_backtrader_{time.strftime('%Y%m%d_%H%M%S')}"
    output_path = os.path.join('backtest_results', run_id)
    generate_backtrader_report(cerebro, results, output_path)
    
    print(f"\n[--- Backtrader run finished! All reports generated in '{output_path}' ---]")